from sqlalchemy import text, create_engine, MetaData
from loguru import logger
import pandas as pd
import numpy as np
import time
import re
import os
from tqdm.notebook import tqdm
import tushare as ts
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

token = '1501ffe708345cffa38d9bbc0bd371e93b4b7412e7a8e1f811d3c442'
ts.set_token(token)
pro = ts.pro_api(token)

dailydata_db_path = f'sqlite:////tushare_database/dailydata/'
minutedata_db_path = f'sqlite:////tushare_database/minutedata/'
tushare_database_folder = 'D:/tushare_database/dailydata'


def 查看_db文件中有哪些表(engine):
    sql = "SELECT name FROM sqlite_master WHERE type='table'"
    rs = engine.execute(sql)
    tables = [table[0] for table in rs.fetchall()]
    return tables


def 查看_db文件中有哪些表2(engine):
    metadata = MetaData()
    metadata.reflect(bind=engine)
    tables = list(metadata.tables.keys())
    return tables


def get_df(期货前缀):
    db_name = f'{期货前缀}.db'
    engine = create_engine(dailydata_db_path + db_name)
    tables = 查看_db文件中有哪些表2(engine)
    dds = []
    for table in tables:
        sql = f"select * from '{table}'"
        dd = pd.read_sql(sql, engine)
        try:
            dd.index = pd.to_datetime(dd['t'])
            dd = dd.drop(['t'], axis=1)
        except:
            dd.index = pd.to_datetime(dd['trade_date'])
            dd = dd.drop(['trade_date'], axis=1)
        dd['code'] = table
        dds.append(dd)
    df = pd.concat(dds)
    df.index.name = 't'
    df = df.reset_index()
    float_col = ['open', 'close', 'low', 'high', 'vol', 'amount']
    df[float_col] = df[float_col].astype(float)
    return df

def get_df2(期货前缀):
    db_name = f'{期货前缀}.db'
    engine = create_engine(dailydata_db_path + db_name)
    tables = 查看_db文件中有哪些表2(engine)
    dds = []
    for table in tables:
        sql = f"select * from '{table}'"
        dd = pd.read_sql(sql, engine)
        try:
            dd.index = pd.to_datetime(dd['t'])
            dd = dd.drop(['t'], axis=1)
        except:
            dd.index = pd.to_datetime(dd['trade_date'])
            dd = dd.drop(['trade_date'], axis=1)
        dd['code'] = table
        dds.append(dd)
    df = pd.concat(dds)
    df.index.name = 't'
    df = df.reset_index()
    float_col = ['open', 'close', 'low', 'high', 'vol', 'amount']
    df[float_col] = df[float_col].astype(float)
    return df


def get_所有期货合约信息():
    交易所s = ['CFFEX', 'DCE', 'CZCE', 'SHFE', 'INE', 'GFEX']
    df_所有合约 = pd.concat([pro.fut_basic(exchange=交易所, fut_type='1') for 交易所 in 交易所s])
    df_所有合约['list_date'] = pd.to_datetime(df_所有合约['list_date'])
    df_所有合约['delist_date'] = pd.to_datetime(df_所有合约['delist_date'])
    df_所有合约 = df_所有合约[df_所有合约['delist_date'] > '2000-01-01']
    df_所有合约 = df_所有合约.sort_values(by=['ts_code']).reset_index(drop=True)
    return df_所有合约


def 更新数据表(需要下载的codes):
    未成功下载的codes = []
    for code in tqdm(需要下载的codes):
        for ss in range(10):
            try:
                table = re.sub("\D", "", code)
                期货前缀 = code.split('.')[0].replace(table, '')
                db_name = f'{期货前缀}.db'
                engine = create_engine(dailydata_db_path + db_name)
                sql = text(f"DROP TABLE IF EXISTS '{table}'")
                engine.execute(sql)
                dd = pro.fut_daily(ts_code=code)
                dd.rename(columns={'trade_date': 't'}, inplace=True)
                dd['t'] = dd['t'].astype(str)
                dd.to_sql(table, engine, index=False, if_exists='append')
                time.sleep(0.3)
                break
            except:
                logger.info('下载出错，重试中..')
                time.sleep(3)
            未成功下载的codes.append(code)
    return 未成功下载的codes


def 更新数据表2(需要下载的codes):
    未成功下载的codes = []
    for code in tqdm(需要下载的codes):
        for ss in range(10):
            try:
                table = re.sub("\D", "", code)
                期货前缀 = code.split('.')[0].replace(table, '')
                db_name = f'{期货前缀}.db'
                engine = create_engine(dailydata_db_path + db_name)
                sql = text(f"DROP TABLE IF EXISTS '{table}'")
                conn = engine.connect()
                conn.execute(sql)
                dd = pro.fut_daily(ts_code=code)
                dd.rename(columns={'trade_date': 't'}, inplace=True)
                dd['t'] = dd['t'].astype(str)
                dd.to_sql(table, engine, index=False, if_exists='append')
                time.sleep(0.3)
                break
            except:
                logger.info('下载出错，重试中..')
                time.sleep(3)
            未成功下载的codes.append(code)
    return 未成功下载的codes


def func_backtest(dff_tp, low_rate, high_rate, ddds, ttp, 期货前缀s, nn):
    if ttp == 1:
        做多的品种 = dff_tp.rank(1, pct=True) > high_rate
        做空的品种 = dff_tp.rank(1, pct=True) < low_rate
    elif ttp == -1:
        做多的品种 = dff_tp.rank(1, pct=True) < low_rate
        做空的品种 = dff_tp.rank(1, pct=True) > high_rate

    dd_rets = []
    for i in range(len(dff_tp) - nn):
        s1_date = dff_tp.index[i]
        s_date = dff_tp.index[i + nn - 1]
        e_date = dff_tp.index[i + nn]
        longs = list(做多的品种.loc[s1_date][做多的品种.loc[s1_date]].index)
        shorts = list(做空的品种.loc[s1_date][做空的品种.loc[s1_date]].index)

        longs_ret = []
        for 品种 in longs:
            longs_ret.append(ddds[期货前缀s.index(品种)][s_date:e_date]['收益率0'])

        shorts_ret = []
        for 品种 in shorts:
            shorts_ret.append(ddds[期货前缀s.index(品种)][s_date:e_date]['收益率0'])

        if len(longs_ret) != 0 and len(shorts_ret) != 0:
            dd_long_ret = pd.concat(longs_ret, axis=1) * 1
            dd_long_ret.columns = longs
            dd_short_ret = pd.concat(shorts_ret, axis=1) * -1
            dd_short_ret.columns = shorts
            #     dd_ret = pd.concat([dd_long_ret,dd_short_ret],axis=1).mean(1)
            dd_ret = (dd_long_ret.mean(1) + dd_short_ret.mean(1)) / 2
            dd_rets.append(dd_ret)
    ddd = pd.concat(dd_rets).sort_index()
    ddd = ddd[~ddd.index.duplicated()]
    return ddd


def get_期货前缀ss(reload):
    if not reload:
        xx = ['A', 'AG', 'AL', 'AP', 'AU', 'B', 'BB', 'BC', 'BU', 'C',
              'CF', 'CJ', 'CS', 'CU', 'CY', 'EB', 'EG', 'ER', 'FB', 'FG',
              'FU', 'HC', 'I', 'IC', 'IF', 'IH', 'IM', 'J', 'JD', 'JM',
              'JR', 'L', 'LH', 'LR', 'LU', 'M', 'MA', 'ME', 'NI', 'NR',
              'OI', 'P', 'PB', 'PF', 'PG', 'PK', 'PM', 'PP', 'RB', 'RI',
              'RM', 'RO', 'RR', 'RS', 'RU', 'SA', 'SC', 'SF', 'SM', 'SN',
              'SP', 'SR', 'SS', 'T', 'TA', 'TC', 'TF', 'TS', 'UR', 'V',
              'WH', 'WR', 'WS', 'WT', 'Y', 'ZC', 'ZN']
        return xx
    期货前缀ss = os.listdir(tushare_database_folder)
    期货前缀ss = [x.split('.')[0] for x in 期货前缀ss]
    期货前缀ss = [x for x in 期货前缀ss if len(x) <= 2]
    期货前缀ss2 = []
    for 期货前缀 in tqdm(期货前缀ss):
        df = get_df(期货前缀)
        df_close = df.pivot(index='t', columns='code', values='close')
        if (df_close.notnull().sum(1) != 0).sum() > 250:
            期货前缀ss2.append(期货前缀)
    return 期货前缀ss2


def get_df_所有合约():
    交易所s = ['CFFEX', 'DCE', 'CZCE', 'SHFE', 'INE', 'GFEX']
    df_所有合约 = pd.concat([pro.fut_basic(exchange=交易所, fut_type='1') for 交易所 in 交易所s])
    df_所有合约['list_date'] = pd.to_datetime(df_所有合约['list_date'])
    df_所有合约['delist_date'] = pd.to_datetime(df_所有合约['delist_date'])
    df_所有合约 = df_所有合约[df_所有合约['delist_date'] > '2010-01-01']
    df_所有合约 = df_所有合约.sort_values(by=['ts_code']).reset_index(drop=True)
    df_所有合约['name'] = df_所有合约['name'].apply(lambda x: re.sub(r'[^\u4e00-\u9fa5]', '', x))
    return df_所有合约


def get_code_delist_date_dict(df_所有合约):
    code_delist_date_dict = dict(zip(df_所有合约['ts_code'].apply(lambda x: x.split('.')[0]),
                                     df_所有合约['delist_date']))
    return code_delist_date_dict


def get_df_合约剩余日期(期货前缀, df_close, code_delist_date_dict):
    last_trade_dates = [code_delist_date_dict.get(期货前缀 + contract) for contract in df_close.columns]
    last_trade_dates = pd.to_datetime(last_trade_dates)
    # 利用了numpy的广播机制
    合约剩余日期 = last_trade_dates.values[None, :] - df_close.index.to_series().values[:, None]
    合约剩余日期 = 合约剩余日期 / np.timedelta64(1, 'D')

    df_合约剩余日期 = pd.DataFrame(合约剩余日期, index=df_close.index, columns=df_close.columns)
    df_合约剩余日期 = df_合约剩余日期.astype(int)
    return df_合约剩余日期


def get_上m期排名第n的col(df_money, m, n):
    return pd.Series(df_money.columns[np.argsort(-df_money.values, axis=1)[:, n]], index=df_money.index).shift(m)


def get_ddd(期货前缀, code_delist_date_dict, pp0, pp1):
    df = get_df(期货前缀)
    df_open = df.pivot(index='t', columns='code', values='open')
    df_high = df.pivot(index='t', columns='code', values='high')
    df_low = df.pivot(index='t', columns='code', values='low')
    df_close = df.pivot(index='t', columns='code', values='close')
    df_close_pct_纵向 = df_close.pct_change(axis=0)
    df_close_pct_横向 = df_close.pct_change(axis=1)
    df_oi = df.pivot(index='t', columns='code', values='oi')
    df_jump = (df_open - df_close.shift(1)) / df_close.shift(-1)
    df_合约剩余日期 = get_df_合约剩余日期(期货前缀, df_close, code_delist_date_dict)
    df_money = df.pivot(index='t', columns='code', values='amount')
    df_money_ratio = df_money.div(df_money.sum(1), axis=0)
    当前代码0 = get_上m期排名第n的col(df_money, 1, 0)
    当前代码1 = get_上m期排名第n的col(df_money, 1, 1)
    当前代码2 = get_上m期排名第n的col(df_money, 1, 2)
    # xx2 = get_所有合约01(期货前缀, 2, pp0, pp1)
    # xx3 = get_所有合约01(期货前缀, 3, pp0, pp1)

    # ***
    ddd月差s = get_ddd月差s(期货前缀)
    dxx2 = get_所有合约0和1(df_money, ddd月差s, 2)
    dxx3 = get_所有合约0和1(df_money, ddd月差s, 3)

    info = {
        '开盘价0': {'df': df_open, 'code': 当前代码0},
        '最高价0': {'df': df_high, 'code': 当前代码0},
        '最低价0': {'df': df_low, 'code': 当前代码0},
        '收盘价0': {'df': df_close, 'code': 当前代码0},
        '收盘价1': {'df': df_close, 'code': 当前代码1},
        '收盘价2': {'df': df_close, 'code': 当前代码2},
        '交易额0': {'df': df_money, 'code': 当前代码0},
        '交易额1': {'df': df_money, 'code': 当前代码1},
        '交易额2': {'df': df_money, 'code': 当前代码2},
        '交易额占比0': {'df': df_money_ratio, 'code': 当前代码0},
        '交易额占比1': {'df': df_money_ratio, 'code': 当前代码1},
        '剩余日0': {'df': df_合约剩余日期, 'code': 当前代码0},
        '剩余日1': {'df': df_合约剩余日期, 'code': 当前代码1},
        '剩余日2': {'df': df_合约剩余日期, 'code': 当前代码2},
        '收益率0': {'df': df_close_pct_纵向, 'code': 当前代码0},
        '收益率1': {'df': df_close_pct_纵向, 'code': 当前代码1},
        '收益率2': {'df': df_close_pct_纵向, 'code': 当前代码2},
        '跳开0': {'df': df_jump, 'code': 当前代码0},
        '跳开1': {'df': df_jump, 'code': 当前代码1},
        '持仓0': {'df': df_oi, 'code': 当前代码0},
        '持仓1': {'df': df_oi, 'code': 当前代码1},
        # '收益率00': {'df': df_close_pct_纵向, 'code': xx2['合约0']},
        # '收益率11': {'df': df_close_pct_纵向, 'code': xx2['合约1']},
        # '收益率000': {'df': df_close_pct_纵向, 'code': xx3['合约0']},
        # '收益率111': {'df': df_close_pct_纵向, 'code': xx3['合约1']},
        # ***
        '收益率x00': {'df': df_close_pct_纵向, 'code': dxx2['合约0']},
        '收益率x11': {'df': df_close_pct_纵向, 'code': dxx2['合约1']},
        '收益率x000': {'df': df_close_pct_纵向, 'code': dxx3['合约0']},
        '收益率x111': {'df': df_close_pct_纵向, 'code': dxx3['合约1']},
    }

    ddd = pd.concat([
        info[name]['df'].apply(lambda x: x.get(info[name]['code'].get(x.name)), axis=1).rename(name)
        for name in info.keys()
    ], axis=1)

    ddd['oi'] = df_oi.sum(1)

    # ddd['合约00'] = xx2['合约0']
    # ddd['合约11'] = xx2['合约1']
    # ddd['合约000'] = xx3['合约0']
    # ddd['合约111'] = xx3['合约1']

    # ***
    ddd['合约x00'] = dxx2['合约0']
    ddd['合约x11'] = dxx2['合约1']
    ddd['合约x000'] = dxx3['合约0']
    ddd['合约x111'] = dxx3['合约1']

    # ***
    ddd['基差率x0011'] = ddd['收益率x00'] - ddd['收益率x11']
    ddd['基差率x000111'] = ddd['收益率x000'] - ddd['收益率x111']

    # ddd['基差率0011'] = ddd['收益率00'] - ddd['收益率11']
    # ddd['基差率000111'] = ddd['收益率000'] - ddd['收益率111']
    ddd['当前代码0'] = 当前代码0
    ddd['当前代码1'] = 当前代码1
    ddd['近月代码'] = np.where(ddd['剩余日0'] < ddd['剩余日1'], 当前代码0, 当前代码1)
    ddd['远月代码'] = np.where(ddd['剩余日0'] > ddd['剩余日1'], 当前代码0, 当前代码1)
    ddd['近月收益率01'] = np.where(ddd['剩余日0'] < ddd['剩余日1'], ddd['收益率0'], ddd['收益率1'])
    ddd['远月收益率01'] = np.where(ddd['剩余日0'] > ddd['剩余日1'], ddd['收益率0'], ddd['收益率1'])
    ddd['近月收益率02'] = np.where(ddd['剩余日0'] < ddd['剩余日2'], ddd['收益率0'], ddd['收益率2'])
    ddd['远月收益率02'] = np.where(ddd['剩余日0'] > ddd['剩余日2'], ddd['收益率0'], ddd['收益率2'])
    ddd['基差率01'] = ddd['近月收益率01'] - ddd['远月收益率01']
    ddd['基差率02'] = ddd['近月收益率02'] - ddd['远月收益率02']
    ddd['展期率01'] = -365 * (ddd['收盘价1'] - ddd['收盘价0']) / (ddd['剩余日1'] - ddd['剩余日0']) / ddd['收盘价0']
    ddd['展期率02'] = -365 * (ddd['收盘价2'] - ddd['收盘价0']) / (ddd['剩余日2'] - ddd['剩余日0']) / ddd['收盘价0']
    ddd['基差RSI'] = -1 * (df_close_pct_横向 * (df_close_pct_横向 > 0)).sum(1) / df_close_pct_横向.abs().sum(1)
    ddd['基差均值'] = -1 * df_close_pct_横向.sum(1) / (df_close.notnull().sum(1) - 1)
    return ddd


def get_因子日度收益(因子, 因子名, 收益率, ttp, thres=0.2):
    tt = 因子.rank(1, pct=True) > 1 - thres
    bb = 因子.rank(1, pct=True) < thres

    做多的品种 = tt if ttp == 1 else bb
    做空的品种 = tt if ttp == -1 else bb
    收益s = []
    for i in range(len(因子) - 1):
        s_date = 因子.index[i]
        e_date = 因子.index[i + 1]
        当期做多 = 做多的品种.loc[s_date].index[做多的品种.loc[s_date]]
        当期做空 = 做空的品种.loc[s_date].index[做空的品种.loc[s_date]]
        做多收益 = 收益率[当期做多][s_date + timedelta(days=1):e_date].mean(1) * 1
        做空收益 = 收益率[当期做空][s_date + timedelta(days=1):e_date].mean(1) * -1
        收益 = (做多收益 + 做空收益) / 2
        # 收益 = 做多收益
        收益s.append(收益)
    rr = pd.concat(收益s)
    rr.name = 因子名
    return rr


def get_因子日度收益2(因子, 因子名, 收益率, ttp, thres=0.2):
    tt = 因子.rank(1) == 3
    bb = 因子.rank(1) == 1

    做多的品种 = tt if ttp == 1 else bb
    做空的品种 = tt if ttp == -1 else bb
    收益s = []
    for i in range(len(因子) - 1):
        s_date = 因子.index[i]
        e_date = 因子.index[i + 1]
        当期做多 = 做多的品种.loc[s_date].index[做多的品种.loc[s_date]]
        当期做空 = 做空的品种.loc[s_date].index[做空的品种.loc[s_date]]
        做多收益 = 收益率[当期做多][s_date + timedelta(days=1):e_date].mean(1) * 1
        做空收益 = 收益率[当期做空][s_date + timedelta(days=1):e_date].mean(1) * -1
        收益 = (做多收益 + 做空收益) / 2
        # 收益 = 做多收益
        收益s.append(收益)
    rr = pd.concat(收益s)
    rr.name = 因子名
    return rr


def add_months_to_code_revised(contract_code, month_diff):
    year, month = int(contract_code[:2]), int(contract_code[2:])
    year += 2000
    start_date = datetime(year, month, 1)
    result_date = start_date + relativedelta(months=+month_diff)
    result_code = f"{result_date.year % 100:02d}{result_date.month:02d}"
    return result_code


def month_difference_str(code1, code2):
    year1, month1 = int(code1[:2]), int(code1[2:])
    year2, month2 = int(code2[:2]), int(code2[2:])

    year1 += 2000
    year2 += 2000

    total_months1 = year1 * 12 + month1
    total_months2 = year2 * 12 + month2
    return total_months1 - total_months2


def get_合约01(given_期货前缀, given_date, nn, pp0, pp1):
    year = given_date.split('-')[0]
    year_1 = year[-2:]
    year_2 = str(int(year) + 1)[-2:]
    month = given_date.split('-')[1]
    month_plus = str(int(month) + 2).zfill(2)
    year_month = year_1 + month
    year_month_plus = add_months_to_code_revised(year_month, nn)

    x0 = pp0[given_期货前缀].copy()
    x1 = pp0[given_期货前缀].copy()
    x0.index = year_1 + x0.index
    x1.index = year_2 + x1.index
    xx = pd.concat([x0, x1])
    合条件合约频次0 = xx[xx.index >= year_month_plus]
    距离0 = pd.Series([month_difference_str(x, year_month) for x in 合条件合约频次0.index.tolist()],
                      index=合条件合约频次0.index)
    分数0 = 合条件合约频次0 / 距离0
    分数0 = 分数0.sort_values(ascending=False)
    合约0 = 分数0.index[0]

    x0 = pp1[given_期货前缀].copy()
    x1 = pp1[given_期货前缀].copy()
    x0.index = year_1 + x0.index
    x1.index = year_2 + x1.index
    xx = pd.concat([x0, x1])
    合条件合约频次1 = xx[xx.index > 合约0]
    距离1 = pd.Series([month_difference_str(x, year_month) for x in 合条件合约频次1.index.tolist()],
                      index=合条件合约频次1.index)
    分数1 = 合条件合约频次1 / 距离1
    分数1 = 分数1.sort_values(ascending=False)
    合约1 = 分数1.index[0]
    return 合约0, 合约1


def get_pp_dict(期货前缀ss):
    pp0 = {}
    pp1 = {}
    for 期货前缀 in tqdm(期货前缀ss):
        df = get_df(期货前缀)
        df_money = df.pivot(index='t', columns='code', values='amount')
        当前代码0 = get_上m期排名第n的col(df_money, 1, 0)
        当前代码1 = get_上m期排名第n的col(df_money, 1, 1)
        pp0.update({期货前缀: 当前代码0.fillna(method='bfill').apply(lambda x: x[-2:]).value_counts()})
        pp1.update({期货前缀: 当前代码1.fillna(method='bfill').apply(lambda x: x[-2:]).value_counts()})
    return pp0, pp1


def get_pp_dict(期货前缀ss):
    pp0 = {}
    pp1 = {}
    for 期货前缀 in tqdm(期货前缀ss):
        df = get_df2(期货前缀)
        df_money = df.pivot(index='t', columns='code', values='amount')
        当前代码0 = get_上m期排名第n的col(df_money, 1, 0)
        当前代码1 = get_上m期排名第n的col(df_money, 1, 1)
        pp0.update({期货前缀: 当前代码0.fillna(method='bfill').apply(lambda x: x[-2:]).value_counts()})
        pp1.update({期货前缀: 当前代码1.fillna(method='bfill').apply(lambda x: x[-2:]).value_counts()})
    return pp0, pp1

def get_所有合约01(期货前缀, nn, pp0, pp1):
    df = get_df(期货前缀)
    df_money = df.pivot(index='t', columns='code', values='amount')
    xx = df_money.index.to_series().apply(lambda x: get_合约01(期货前缀, str(x).split()[0], nn, pp0, pp1))
    xx = xx.apply(pd.Series)
    xx.columns = ['合约0', '合约1']
    return xx


def get_ddd月差s(期货前缀):
    df = get_df(期货前缀)
    df_money = df.pivot(index='t', columns='code', values='amount')
    dd = df_money.resample('M').sum()
    ddd月差s = {}
    for the_date in dd.index:
        year, month = str(the_date).split('-')[:2]
        the_date_month_code = year[-2:] + month
        ddd = dd.loc[the_date]
        ddd.index = [month_difference_str(x, the_date_month_code) for x in ddd.index]
        if not ddd月差s.get(month):
            ddd月差s.update({month: []})
        ddd月差s[month].append(ddd)
    return ddd月差s


def get_合约0和1(current_month_code, ddd月差s, 最小月差距离):
    month = current_month_code[-2:]
    dx = pd.concat(ddd月差s[month], axis=1).mean(1).sort_index()
    dx = dx[dx.index >= 最小月差距离].sort_values(ascending=False).iloc[:2].sort_index()
    月差0 = dx.index[0]
    月差1 = dx.index[1]
    合约0 = add_months_to_code_revised(current_month_code, 月差0)
    合约1 = add_months_to_code_revised(current_month_code, 月差1)
    return 合约0, 合约1


def get_month_code(x):
    x = str(x)
    year, month = x.split('-')[:2]
    return year[-2:] + month


def get_所有合约0和1(df_money, ddd月差s, nn):
    xx = df_money.index.to_series().apply(lambda x: get_合约0和1(get_month_code(str(x).split()[0]), ddd月差s, nn))
    xx = xx.apply(pd.Series)
    xx.columns = ['合约0', '合约1']
    return xx
